from flask import Flask, request, jsonify
from gmssl import func, sm2
from datetime import datetime, timedelta
import json
import base64

app = Flask(__name__)

class SM2CertificateAuthority:
    def __init__(self):
        # 生成 CA 的 SM2 密钥对
        self.private_key = func.random_hex(32)  # 32字节私钥 (64字符 hex)
        self.public_key = None
        self.crypt = None
        self._generate_keys()

    def _generate_keys(self):
        # 生成私钥（64字符 hex）
        self.private_key = func.random_hex(32)  # 32字节 → 64字符 hex

        # 创建 CryptSM2 实例（不传公钥）
        self.crypt = sm2.CryptSM2(public_key=None, private_key=self.private_key)

        # 触发公钥生成：执行一次 dummy 签名（不需要真的发送）
        # 任意数据签名会触发公钥计算
        dummy_data = "00"  # 任意数据
        _ = self.crypt.sign_with_sm3(dummy_data, func.random_hex(32))

        # 签名后，gmssl 会自动计算公钥并存入 self.crypt.public_key
        if self.crypt.public_key is None:
           raise RuntimeError("Failed to generate SM2 public key from private key")

        # 获取公钥（已经是 hex 字符串，未压缩，以 '04' 开头）
        self.public_key = self.crypt.public_key

        # 可选：去掉 '04' 前缀（如果你希望存储为纯坐标）
        # self.public_key_no_prefix = self.public_key[2:] if self.public_key.startswith('04') else self.public_key

        print(f"CA 私钥: {self.private_key}")
        print(f"CA 公钥: {self.public_key}")

    def _sign_data(self, data: str) -> str:
        """使用 SM2 对数据进行签名，返回 hex 签名"""
        return self.crypt.sign_with_sm3(data, func.random_hex(32))

    def _verify_signature(self, data: str, signature: str, public_key: str) -> bool:
        """验证 SM2 签名"""
        verifier = sm2.CryptSM2(public_key=public_key, private_key=None)
        return verifier.verify_with_sm3(signature, data)

    def get_root_pubkey(self):
        """返回 CA 的公钥信息（可视为“根证书”）"""
        return {
            "ca": "SM2 CA Authority",
            "public_key": self.public_key,
            "algorithm": "SM2",
            "timestamp": datetime.utcnow().isoformat() + "Z"
        }

    def issue_certificate(self, subject: str, subject_public_key: str, validity_days=365):
        """
        模拟签发 SM2 证书
        :param subject: 主题（如 common_name）
        :param subject_public_key: 用户的 SM2 公钥（hex）
        :param validity_days: 有效期
        :return: 包含签名的“证书”
        """
        # 构造证书内容
        cert_data = {
            "version": "SM2-v1",
            "serial_number": func.random_hex(16),  # 模拟序列号
            "subject": subject,
            "subject_public_key": subject_public_key,
            "issuer": "SM2 CA Authority",
            "not_before": datetime.utcnow().strftime("%Y-%m-%d %H:%M:%S"),
            "not_after": (datetime.utcnow() + timedelta(days=validity_days)).strftime("%Y-%m-%d %H:%M:%S"),
            "algorithm": "SM2-SM3"
        }

        # 将证书内容转为字符串用于签名
        data_to_sign = json.dumps(cert_data, sort_keys=True, separators=(',', ':'))

        # 使用 CA 的 SM2 私钥签名
        signature = self._sign_data(data_to_sign)

        # 返回“证书”：数据 + 签名 + 算法标识
        certificate = {
            "data": cert_data,
            "signature": signature,
            "signature_algorithm": "SM2-SM3",
            "ca_public_key": self.public_key  # 便于验证
        }

        return certificate

    def verify_certificate(self, certificate):
        """验证证书签名是否由 CA 签发"""
        cert_data = certificate["data"]
        signature = certificate["signature"]
        ca_pubkey = certificate.get("ca_public_key") or self.public_key

        data_to_sign = json.dumps(cert_data, sort_keys=True, separators=(',', ':'))
        return self._verify_signature(data_to_sign, signature, ca_pubkey)


# 初始化 CA
ca = SM2CertificateAuthority()


@app.route('/ca_root', methods=['GET'])
def get_ca_root():
    """返回 CA 的公钥（相当于根证书）"""
    root_info = ca.get_root_pubkey()
    return jsonify(root_info), 200


@app.route('/request_cert', methods=['POST'])
def request_certificate():
    """接收请求并颁发 SM2 证书"""
    data = request.json
    if not data or 'subject' not in data or 'public_key' not in data:
        return jsonify({'error': 'Missing subject or public_key in request'}), 400

    subject = data['subject']
    subject_public_key = data['public_key']  # hex 格式的 SM2 公钥
    validity_days = data.get('validity_days', 365)

    # 验证公钥格式（简单检查）
    if not (len(subject_public_key) == 130 or len(subject_public_key) == 128) or not subject_public_key.startswith('04'):
        # 未压缩公钥应为 04 + x(64) + y(64) => 130 字符 hex
        if len(subject_public_key) == 128:
            subject_public_key = '04' + subject_public_key  # 自动补全
        else:
            return jsonify({'error': 'Invalid SM2 public key format'}), 400

    try:
        certificate = ca.issue_certificate(subject, subject_public_key, validity_days)
        return jsonify({
            'certificate': certificate,
            'message': 'Certificate issued with SM2-SM3'
        }), 200
    except Exception as e:
        return jsonify({'error': str(e)}), 500


@app.route('/verify_cert', methods=['POST'])
def verify_certificate():
    """验证一个 SM2 证书是否由本 CA 签发"""
    data = request.json
    if not data or 'certificate' not in data:
        return jsonify({'error': 'Missing certificate'}), 400

    cert = data['certificate']
    try:
        is_valid = ca.verify_certificate(cert)
        return jsonify({'valid': is_valid}), 200
    except Exception as e:
        return jsonify({'error': str(e)}), 400


if __name__ == '__main__':
    print("Starting SM2 CA Server on https://0.0.0.0:28291")
    app.run(host='0.0.0.0', port=28291, ssl_context='adhoc')